package goComfyuiApi

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
)

type OutputInfo struct {
	Images []OutputImageInfo `json:"images,omitempty"`
	Text   []string          `json:"text,omitempty"`
}

type historyStatusInfo struct {
	StatusStr string `json:"status_str"`
	Completed bool   `json:"completed"`
}

type historyTaskInfo struct {
	Status  historyStatusInfo      `json:"status"`
	Outputs map[string]*OutputInfo `json:"outputs"`
}

type historyResult map[string]*historyTaskInfo // map[promptID][_]

func getWorkflowOutput(comfyuiUrl, promptID string) (*historyTaskInfo, error) {
	url := genComfyUIApiUrl(comfyuiUrl, fmt.Sprintf("history/%s", promptID))
	resp, err := http.Get(url)
	if err != nil {
		return nil, fmt.Errorf("failed to get history: %v", err)
	}
	defer resp.Body.Close()
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("failed to read history response: %v", err)
	}
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("history request failed: %s", string(body))
	}

	var historyResponse historyResult
	err = json.Unmarshal(body, &historyResponse)
	if err != nil {
		return nil, fmt.Errorf("failed to decode history response: %v", err)
	}

	taskInfo, ok := historyResponse[promptID]
	if !ok {
		return nil, nil //fmt.Errorf("promptID not found: %s", promptID)
	}

	return taskInfo, nil
}

// 获取执行结果中的图片
func ApiGetWorkflowOutputImages(comfyuiUrl, promptID string) ([]OutputImageInfo, error) {
	outImages := []OutputImageInfo{}

	taskInfo, err := getWorkflowOutput(comfyuiUrl, promptID)
	if err != nil {
		return outImages, err
	}

	if taskInfo == nil {
		return outImages, nil
	}

	outputs := taskInfo.Outputs
	// 遍历outputs
	for _, o := range outputs {

		// 遍历images
		for _, img := range o.Images {
			if img.Type != "output" {
				continue
			}
			outImages = append(outImages, img)
		}

	}

	return outImages, nil
}
